# agent_structured_advanced_demo.py：多步工具调用 + 参数校验 + 中文提示词 + JSON 调用

from langchain.tools import tool
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
import os
import json
import re
from typing import Dict, Any, List, Tuple


# =============== 环境与模型 ===============
load_dotenv()

llm = ChatOpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="qwen-turbo"
)


# =============== 定义工具（Tools） ===============
@tool(description="用于查询指定城市的当前天气。参数：city（必填，城市名，示例：北京、上海），返回值：温度和天气状况。用户问城市天气必须调用此工具，不可编造数据。")
def WeatherTool(city: str) -> str:
    weather_data = {
        "北京": "温度18-25℃，晴，微风",
        "上海": "温度20-28℃，多云，东南风3级",
        "广州": "温度25-32℃，阴，有小雨"
    }
    return weather_data.get(city, f"暂未获取到{city}的天气数据")


@tool(description="用于计算加减乘除运算。参数：expression（必填，数学表达式，示例：2+3*4、10-5，无空格），返回值：计算结果。仅处理数字和+、-、*、/。")
def CalculatorTool(expression: str) -> str:
    try:
        result = eval(expression)
        return f"计算结果：{expression} = {result}"
    except Exception as e:
        return f"计算失败：{str(e)}，请检查表达式（示例：2+3，不要有空格）"


@tool(description="用于查询公司内部知识库（如年假政策、产品价格）。参数：question（必填，查询问题，示例：公司年假政策、产品A价格），返回值：知识库答案。公司相关问题必须调用此工具，不可编造。")
def RAGTool(question: str) -> str:
    rag_data = {
        "公司年假政策": "员工入职满1年5天年假，满3年10天，满5年15天。",
        "产品A的价格": "产品A基础版999元/年，高级版1999元/年，企业版需定制。",
        "加班补贴标准": "工作日20元/小时，周末30元/小时，法定节假日50元/小时。"
    }
    for key, value in rag_data.items():
        if key in question:
            return value
    return f"知识库中未找到「{question}」的信息，请联系相关部门。"


# =============== 工具注册与参数校验 Schema ===============
REGISTERED_TOOLS: Dict[str, Any] = {
    WeatherTool.name: WeatherTool,
    CalculatorTool.name: CalculatorTool,
    RAGTool.name: RAGTool,
}

# 简单参数 Schema（演示用途）：
# - required: 必填字段列表
# - types: 字段类型（仅演示 str、int、float、bool）
TOOL_SCHEMAS: Dict[str, Dict[str, Any]] = {
    WeatherTool.name: {
        "required": ["city"],
        "types": {"city": str},
    },
    CalculatorTool.name: {
        "required": ["expression"],
        "types": {"expression": str},
    },
    RAGTool.name: {
        "required": ["question"],
        "types": {"question": str},
    },
}


def validate_parameters(tool_name: str, params: Dict[str, Any]) -> Tuple[bool, str]:
    schema = TOOL_SCHEMAS.get(tool_name)
    if not schema:
        return True, ""
    # 检查必填
    for field in schema.get("required", []):
        if field not in params:
            return False, f"缺少必填参数: {field}"
    # 检查类型
    types_map = schema.get("types", {})
    for k, expected_type in types_map.items():
        if k in params and not isinstance(params[k], expected_type):
            return False, f"参数类型错误: {k} 期望 {expected_type.__name__}, 实际 {type(params[k]).__name__}"
    return True, ""


# =============== 中文提示词：要求 JSON 工具调用，多步循环 ===============
SYSTEM_RULES = (
    "你是一个严格遵守规则的中文智能Agent。你可以使用以下工具来完成任务：\n"
    "{tools}\n\n"
    "必须遵守以下调用规则：\n"
    "1) 若无需工具即可回答，请直接输出最终答案，且不要输出任何工具调用。\n"
    "2) 若需要调用工具，请严格输出一个 JSON 代码块，格式如下：\n"
    "```json\n"
    "{{\"name\": \"工具名\", \"parameters\": {{\"参数名\": \"参数值\"}}}}\n"
    "```\n"
    "3) JSON 代码块之外不要输出任何多余文字。\n"
    "4) 工具返回结果后，如果还需要调用其它工具，请继续按相同格式输出下一个 JSON；如果已经能回答，请输出中文最终答案，不要再输出 JSON。\n"
)


def render_tools_doc() -> str:
    lines: List[str] = []
    for tool_obj in REGISTERED_TOOLS.values():
        lines.append(f"- {tool_obj.name}{tool_obj.args} - {tool_obj.description}")
    return "\n".join(lines)


def build_step_prompt(user_input: str, transcript: str) -> str:
    tools_doc = render_tools_doc()
    parts = [
        SYSTEM_RULES.format(tools=tools_doc),
        "用户问题：" + user_input,
    ]
    if transcript:
        parts.append("已知上下文（逐步记录）：\n" + transcript.strip())
    parts.append("请按规则作答。")
    return "\n".join(parts)


# =============== JSON 提取与执行 ===============
JSON_BLOCK_RE = re.compile(r"```json\s*(\{[\s\S]*?\})\s*```", re.IGNORECASE)


def try_extract_json_block(text: str) -> Dict[str, Any]:
    match = JSON_BLOCK_RE.search(text)
    if not match:
        text = text.strip()
        if text.startswith("{") and text.endswith("}"):
            try:
                return json.loads(text)
            except Exception:
                return {}
        return {}
    try:
        return json.loads(match.group(1))
    except Exception:
        return {}


def call_tool_json(tool_call: Dict[str, Any]) -> Tuple[bool, str]:
    name = tool_call.get("name")
    params = tool_call.get("parameters", {})
    tool_fn = REGISTERED_TOOLS.get(name)
    if tool_fn is None:
        return False, f"未找到工具：{name}"

    # 参数校验
    ok, err = validate_parameters(name, params)
    if not ok:
        return False, f"参数校验失败：{err}"

    print("\n====== Tool JSON 调用 ======")
    print(json.dumps({"name": name, "parameters": params}, ensure_ascii=False, indent=2))

    # 优先使用标准调用接口
    try:
        if hasattr(tool_fn, "invoke"):
            result = tool_fn.invoke(params)
        elif hasattr(tool_fn, "_run"):
            result = tool_fn._run(**params)
        else:
            # 尝试直接以关键字参数方式调用原始函数（仅当可行时）
            result = tool_fn(**params)
        return True, result
    except TypeError as e:
        return False, f"工具调用方式错误：{str(e)}"
    except Exception as e:
        return False, f"工具执行异常：{str(e)}"


# =============== 多步控制循环 ===============
def run_multi_step(user_input: str, max_steps: int = 5) -> str:
    transcript = ""
    for step in range(1, max_steps + 1):
        prompt = build_step_prompt(user_input, transcript)
        print("\n====== 第 %d 步：发送给模型的提示词 ======" % step)
        print(prompt)

        response = llm.invoke(prompt)
        content = response.content if hasattr(response, "content") else str(response)

        tool_call = try_extract_json_block(content)

        # 若没有 JSON，视为模型给出最终答案
        if not tool_call:
            final_answer = content.strip()
            print("\n====== 最终答案 ======")
            print(final_answer)
            return final_answer

        # 有 JSON：执行工具
        ok, tool_result = call_tool_json(tool_call)
        print("\n====== 工具返回结果 ======")
        print(tool_result)

        # 更新逐步记录（便于下一轮继续思考）
        transcript += (
            f"\n[Step {step}] 工具调用: {json.dumps(tool_call, ensure_ascii=False)}\n"
            f"[Step {step}] 工具结果: {tool_result}\n"
        )

        # 把“参数校验失败/执行失败”显式反馈给模型，引导其修正 JSON
        if not ok:
            correction_prompt = (
                "上一步工具调用存在问题（参数或执行失败），请根据错误信息修正你的 JSON 并重新调用工具。\n"
                f"错误信息：{tool_result}\n"
                "只输出 JSON，不要输出其它内容。"
            )
            transcript += "[Assistant 提示] " + correction_prompt + "\n"
            continue

        # 工具成功：提示可以继续选择下一步或给出最终答案
        continue_prompt = (
            "根据上述工具结果，若已能完成任务，请输出中文最终答案（不要再输出 JSON）；"
            "若仍需继续调用其他工具，请再次按规则输出 JSON。"
        )
        transcript += "[Assistant 提示] " + continue_prompt + "\n"

    # 超过最大步数
    return "已达到最大思考步数，但仍未给出最终答案。请缩小问题范围或提高步数。"


if __name__ == "__main__":
    print("场景A：无需工具 → 什么是人工智能？")
    out_a = run_multi_step("什么是人工智能？用简单的话解释。", max_steps=3)
    print("答案：", out_a)

    print("\n" + "=" * 50)
    print("场景B：单工具 → 北京天气")
    out_b = run_multi_step("查询北京今天的天气。", max_steps=3)
    print("答案：", out_b)

    print("\n" + "=" * 50)
    print("场景C：多步链条 → 先查上海天气，再计算预算 150+200+300")
    out_c = run_multi_step("查上海今天的天气，再计算一套穿搭预算（150+200+300）。", max_steps=5)
    print("答案：", out_c)


